# 用于直观地显示出图像中识别到的轮廓

import cv2
import matplotlib.pyplot as plt
import numpy as np

# 定义红色的两个HSV范围
red_lower1 = np.array([0, 80, 236])
red_upper1 = np.array([10, 255, 255])
red_lower2 = np.array([150, 80, 236])
red_upper2 = np.array([190, 255, 255])

# 定义黄色的HSV范围
yellow_lower = np.array([20, 80, 236])
yellow_upper = np.array([30, 255, 255])

kernel = np.ones((3, 3), np.uint8)

# width, height = pyautogui.size()
# region = (0, 0, width, height)
# screenshot = pyautogui.screenshot(region=region)

bgr_img = cv2.imread("image.png", cv2.IMREAD_COLOR)
screenshot = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)


def find_yellow_rectangles(_hsv_img):
    mask = cv2.inRange(_hsv_img, yellow_lower, yellow_upper)
    find_rectangles(mask)


def find_red_rectangles(_hsv_img):
    mask_red1 = cv2.inRange(_hsv_img, red_lower1, red_upper1)
    mask_red2 = cv2.inRange(_hsv_img, red_lower2, red_upper2)
    mask_red_combined = cv2.bitwise_or(mask_red1, mask_red2)
    find_rectangles(mask_red_combined)


def find_rectangles(mask):
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 在原图上绘制轮廓
    image_with_contours = draw_contours_on_image(screenshot, contours)
    # 显示图像
    plt.figure(figsize=(16, 9), dpi=150)
    plt.imshow(image_with_contours)
    plt.axis('off')  # 关闭坐标轴
    plt.show()
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        print("w:", w, "h:", h)


def draw_contours_on_image(image, contours, color=(0, 255, 0), thickness=1):
    # 创建原图的一个副本，以避免修改原始图像
    image_with_contours = image.copy()

    # 使用cv2.drawContours绘制轮廓到图像副本上
    cv2.drawContours(image_with_contours, contours, -1, color, thickness)

    return image_with_contours


hsv_img = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2HSV)
# find_yellow_rectangles(hsv_img)
find_red_rectangles(hsv_img)
